import libvirt
import libxml2
import logging
import os
import os.path
import subprocess
import sys

from urlparse import urlsplit

class VirtManager:
	'''
	This class is a convenience class around a libvirt domain,
	it contains a couple of methods for interacting with the domain.
	
	Users must remember to call start_domain to start the domain and call shutdown when finished
	'''
	def __init__(self, config_path, uri=None, log=None):
		self.config_path = config_path
		if uri:
			self.__get_remote(uri)
		
		self.name, xml = self.__read_xml(config_path)
		if log:
			self.log = log
		else:
			self.log = logging.getLogger(self.name)
		self.log.debug('Connecting to uri: %s', uri)
		self.connection = libvirt.open(uri) # None defaults to the first hypervisor libvirt can find
		
		self.log.debug('Domain name is "%s"', self.name)
		try:
			self.__connect_domain(self.name, xml)
		except:
			#Be sure to cleanup in case of emergency, then re-raise error
			self.connection.close()
			raise
		
		#Prepare stream and console in case user wants to write to VM console
		self.stream = None
		self.console = None
		
	def __repr__(self):
		ret = "%s(config_path=%r, uri=%r, log=%r), " %(self.__class__.__name__, self.config_path, self.connection.getURI(), self.log)
		ret += "Domain name is {name}, ".format(name=self.name)
		ret += "Connection %r, " % self.connection
		ret += "Running on remote host '{ip}', ".format(ip=self.remote_ip) if self.remote else "Running on local host, "
		ret += "Remote username is '{user}', ".format(user=self.remote_usr) if self.remote and self.remote_usr != None else "Using default username, "
		ret += "Stream is %r, " % self.stream
		ret += "Console is %r" % self.console
		return ret
		
	def __get_remote(self, uri):
		'''
		Helper method to extract url and user if this is a remote domain
		'''
		uri_match = urlsplit(uri)
		if uri_match.netloc:
			self.remote = True
			try:
				(self.remote_usr, self.remote_ip) = uri_match.netloc.split('@')
			except ValueError:
				#Could not unpack into two values, means that
				#there is no user in url
				self.remote_ip = uri_match.netloc
		else:
			self.remote = False

	def __read_xml(self, config_path):
		'''
		Read the XML configuration in the file defined in config_path
		returns a tuple (name, xml)
		'''
		with open(config_path, 'r') as f:
			xmldesc = f.read()
		doc = libxml2.parseDoc(xmldesc)
		name = doc.xpathNewContext().xpathEval("/domain/name")[0].content
		#self.log.debug('Read XML description of domain')
		return (name, xmldesc)
	
	def __connect_domain(self, name, xml=None):
		'''
		Connect to the domain with name equal to 'name', if 'name' is not a
		defined domain, the method expects 'xml' to contain a valid xml
		description of a domain.
		
		This method will throw an exception if it can't create the domain
		'''
		try:
			self.log.debug('Looking up domain by name')
			self.domain = self.connection.lookupByName(self.name)
		except libvirt.libvirtError:
			self.log.debug('Domain is not defined, defining from XML')
			self.domain = self.connection.defineXML(xml)
		if self.domain == None:
			self.log.critical('Could not connect to domain %s', name)
			raise RuntimeError('Could not connect to domain with name {nm}'.format(nm=name))
		
	def start_domain(self):
		'''
		Start the domain, this method will throw an exception if the domain is runnin.
		'''
		try:
			self.log.debug('Starting domain')
			self.domain.create()
		except libvirt.libvirtError, err:
			#This could be caused by the domain already running, but it could also
			#be a problem, let the caller decide what to do
			self.log.debug('Got exception trying to start domain\n%s', err)
			raise
			
	def __write_console(self, toWrite):
		'''
		Write the specified string to the console of the domain
		
		This method expects the Stream and Console to have been properly opened
		before a call to this method.
		'''
		#Nothing is stripped from the actual command, but for printability we remove '\n'
		self.log.debug('Writing "%s" to console', toWrite.strip('\n'))
		self.stream.send(toWrite)
		
	def run_command(self, command):
		'''
		Run the specified command in the console of the domain.
		The command will run asynchronous.
		The command is assumed to be a list of arguments to be passed. Eg. ['python', 'test.py']
		'''
		self.__start_console()
		self.log.debug('Got command %s', command)
		#Remove whitespaces from arguments
		c = " ".join(map(str.strip, command))
		self.log.debug('Running command: "%s"', c)
		#For safety we add a newline to the start of the command
		#so we are guaranteed that the command will run
		self.__write_console('\n{}\n'.format(c))
		
	def __create_folder_local(self, local_dest):
		'''
		Create a folder on the local machine. 
		Works on both Windows and Unix
		'''
		try:
			self.log.debug('Creating "%s" on local machine', local_dest)
			os.mkdir(local_dest)
		except OSError:
			pass #Directory is already present
		
	def __create_folder_host(self, domain_dest, windows=False):
		'''
		Create a folder on the host domain, method work both local and remote.
		Pass in windows=True to have 'domain_dest' converted to Windows path according to
		ntpath.abspath() in Python.
		Works on Windows and Unix
		'''
		if windows:
			import ntpath
			domain_dest = ntpath.abspath(domain_dest)
		self.log.debug('Creating "%s" on the remote host %s', domain_dest, self.remote_ip)
		self.run_command(['mkdir', domain_dest])
		
	def upload_folder(self, local_src, domain_dest):
		'''
		Convenience method to upload a folder to the VM
		'''
		self.__create_folder_host(self, domain_dest)
		self.log.debug('Uploading folder "%s" to destination "%s"', local_src, domain_dest)
		for item in os.listdir(local_src):
			full_path = os.path.abspath(local_src + item)
			if os.path.isfile(full_path):
				self.upload_file(full_path)
			elif os.path.isfolder(full_path):
				self.upload_folder(full_path, domain_dest)
			else:
				self.log.warn('Could not upload "%s", neither file nor folder', full_path)
		
	def download_folder(self, domain_src, local_dest):
		'''
		Convenience method to download a folder from the VM
		'''
		self.__create_folder_local(self, local_dest)
		
		if remote:
			if self.remote_usr:
				rem = '{user}@{host}:{src}'.format(user=self.remote_usr, host=self.remote_ip, src=domain_src)
			else:
				rem = '{host}:{src}'.format(host=self.remote_ip, src=domain_src)
			ret_code = subprocess.check_call(['scp', '-r', rem, local_dest])
		else:
			ret_code = subprocess.check_call(['cp', '-r', domain_src, local_dest])
		if ret_code:
			self.log.error('Problem when downloading folder from domain, return code was: "%s"', ret_code)
			
	def upload_file(self, local_src, domain_dest):
		'''
		Upload the 'local_src' file to 'domain_dest'
		'''
		self.log.debug('Uploading "%s" to "%s" on domain', local_src, domain_dest)
		if self.remote:
			self.__remote_upload(local_src, domain_dest)
		else:
			self.__local_upload(local_src, domain_dest)
		
	def download_file(self, domain_src, local_dest):
		'''
		Download a file at 'domian_src' to destination 'local_dest'
		'''
		self.log.debug('Downloading "%s" from domain to "%s"', domain_src, local_dest)
		if self.remote:
			self.__remote_download(domain_src, local_dest)
		else:
			self.__local_download(domain_src, local_dest)
		
	def __remote_upload(self, local_src, domain_dest):
		'''
		Upload file to VM on remote host
		'''
		if self.remote_usr:
			rem = '{user}@{host}:{src}'.format(user=self.remote_usr, host=self.remote_ip, src=domain_dest)
		else:
			rem = '{host}:{src}'.format(host=self.remote_ip, src=domain_dest)
		ret_code = subprocess.check_call(['scp', local_src, rem])
		if ret_code:
			self.log.error('SCP returned "%s" when exiting!', ret_code)
		
	def __local_upload(self, local_src, domain_dest):
		'''
		Upload file to VM on local host
		'''
		ret_code = subprocess.check_call(['cp', local_src, domain_dest])
		if ret_code:
			self.log.error('Copy command returned "%s"!', ret_code)
		
	def __remote_download(self, domain_src, local_dest):
		'''
		Download file from VM on remote host
		'''
		if self.remote_usr:
			rem = '{user}@{host}:{src}'.format(user=self.remote_usr, host=self.remote_ip, src=domain_src)
		else:
			rem = '{host}:{src}'.format(host=self.remote_ip, src=domain_src)
		ret_code = subprocess.check_call(['scp', rem, local_dest])
		if ret_code:
			self.log.error('SCP returned "%s" when exiting!', ret_code)
		
	def __local_download(self, domain_src, local_dest):
		'''
		Download file from VM on local host
		'''
		ret_code = subprocess.check_call(['cp', domain_src, local_dest])
		if ret_code:
			self.log.error('Copy command returned "%s"!', ret_code)
		
	def __start_console(self):
		'''
		Helper method used to start a console connection to the domain. Methods which deppend
		on a console connection should call this method before starting to use the console.
		'''
		if self.stream == None:
			self.log.debug('Creating stream to use with console')
			self.stream = self.connection.newStream(0)
		if self.console == None:
			self.log.debug('Opening console')
			self.console = self.domain.openConsole(None, self.stream, 0)
		
	def shutdown(self, undefine=False):
		'''
		Close all open libvirt connections
		
		If undefine is true the domain xml will also be removed. Lookup by name
		will fail after undefine.
		
		Shutdown might throw an exception, but will continue
		on its way. This is done to ensure that the shutdown is completed properly
		and all resources regarding libvirt is closed.
		'''
		self.log.debug('Starting to shutdown')
		
		try:
			self.stream.finish()
		except AttributeError:
			pass #If stream was not used it will be None
		except libvirt.libvirtError as err:
			error = err
			self.log.exception('Got exception trying to finish stream')
		
		try:
			self.domain.shutdown()
		except libvirt.libvirtError as err:
			if self.connection.getType().lower() in ['lxc']:
				pass #LXC driver does not support shutdown
			else:
				error = err
				self.log.exception('Got exception trying to shutdown domain')
		finally:
			try:
				self.domain.destroy()
			except Exception as err:
				error = err
			
		if undefine:
			self.log.debug('Undefining domain')
			try:
				self.domain.undefine()
			except libvirt.libvirtError as err:
				error = err
				self.log.exception('Got exception while trying to undefine domain')
		
		try:
			self.connection.close()
		except libvirt.libvirtError as err:
			error = err
			self.log.exception('Got exception trying to close libvirt connection')
		
		if error:
			#This is not a pretty solution, but in case of a problem
			#we will shutdown the other connections to libvirt and
			#raise the last error.
			raise error
		
		self.log.debug('Shutdown complete')
		
		
def main():
	logging.basicConfig(level=logging.DEBUG)
	vmm = VirtManager('configs/lxc-1.xml', 'lxc:///')
	vmm.start_domain()
	try:
		vmm.run_command(['ls', '-l'])
	finally:
		vmm.shutdown()
if __name__ == '__main__':
	main()